//------------------------------------------------------------------------------
// File: GrabBitmaps.cpp
//
// Desc: DirectShow sample code - GrabBitmaps sample
//       This console app will open a long AVI file in the parent directory,
//       create a filter graph with a sample grabber filter,
//       read frames out of it every second for a few seconds, 
//       and write the frames to BMP files in the current directory.
//
// Copyright (c) Microsoft Corporation.  All rights reserved.
//------------------------------------------------------------------------------

#include <windows.h>
#include <streams.h>
#include <stdio.h>
#include <atlbase.h>
#include <qedit.h>
#include <ximage.h>             // please refer to http://www.codeproject.com/bitmap/cximage.asp to get it.
#include <time.h>

// Function prototypes
int GrabBitmaps(TCHAR * lpszInputFile, TCHAR *lpszOutputFile, int nOutWidth, int nOutHeight, int nImageCounter );
HRESULT GetPin(IBaseFilter * pFilter, PIN_DIRECTION dirrequired,  int iNum, IPin **ppPin);
IPin *  GetInPin ( IBaseFilter *pFilter, int Num );
IPin *  GetOutPin( IBaseFilter *pFilter, int Num );

// Constants
#define NUM_FRAMES_TO_GRAB  20


//
// Implementation of CSampleGrabberCB object
//
// Note: this object is a SEMI-COM object, and can only be created statically.

class CSampleGrabberCB : public ISampleGrabberCB 
{

public:

    CSampleGrabberCB( TCHAR* _lpszOutputFile, int _nOutWidth, int _nOutHeight, int _nImageCounter )
    {
        _tcscpy( lpszOutFile, _lpszOutputFile );
        nOutWidth  = _nOutWidth;
        nOutHeight = _nOutHeight;
        fPrevlDiffSquare = FLT_MAX;
        nImageCounter = _nImageCounter;
        nImageGened   = 0;
    }
    
    TCHAR lpszOutFile[MAX_PATH];
    int nOutWidth;
    int nOutHeight;
    int nImageCounter;
    int nImageGened;

    // These will get set by the main thread below. We need to
    // know this in order to write out the bmp
    long Width;
    long Height;

    // Fake out any COM ref counting
    //
    STDMETHODIMP_(ULONG) AddRef() { return 2; }
    STDMETHODIMP_(ULONG) Release() { return 1; }

    // Fake out any COM QI'ing
    //
    STDMETHODIMP QueryInterface(REFIID riid, void ** ppv)
    {
        CheckPointer(ppv,E_POINTER);
        
        if( riid == IID_ISampleGrabberCB || riid == IID_IUnknown ) 
        {
            *ppv = (void *) static_cast<ISampleGrabberCB*> ( this );
            return NOERROR;
        }    

        return E_NOINTERFACE;
    }


    // We don't implement this one
    //
    STDMETHODIMP SampleCB( double SampleTime, IMediaSample * pSample )
    {
        return 0;
    }


    // The sample grabber is calling us back on its deliver thread.
    // This is NOT the main app thread!
    //
    STDMETHODIMP BufferCB( double SampleTime, 
                           BYTE * pBuffer, 
                           long BufferSize )
    {
        int nBmpBufferSize = (int)sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + BufferSize;
        BYTE *pbBMP = new BYTE[ nBmpBufferSize ];

        BITMAPFILEHEADER *pbfh = (BITMAPFILEHEADER*)pbBMP;
        memset( pbfh, 0, sizeof( BITMAPFILEHEADER ) );
        pbfh->bfType = 'MB';
        pbfh->bfSize = sizeof( BITMAPFILEHEADER ) + BufferSize + sizeof( BITMAPINFOHEADER );
        pbfh->bfOffBits = sizeof( BITMAPINFOHEADER ) + sizeof( BITMAPFILEHEADER );


        BITMAPINFOHEADER *pbih = (BITMAPINFOHEADER *)(pbBMP + sizeof(BITMAPFILEHEADER)) ;
        memset( pbih, 0, sizeof( BITMAPINFOHEADER ) );
        pbih->biSize = sizeof( BITMAPINFOHEADER);
        pbih->biWidth = Width;
        pbih->biHeight = Height;
        pbih->biPlanes = 1;
        pbih->biBitCount = 24;

        BYTE *pBites = pbBMP + (int)(sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER));
        
        memcpy( pBites, pBuffer, BufferSize );

        CxImage  image( pbBMP, nBmpBufferSize, CXIMAGE_FORMAT_BMP );
        if (image.IsValid())
        {
            float fCurDifSquare = GetGrayPixelSize( image );
            if( fPrevlDiffSquare > fCurDifSquare 
                || nImageCounter - 1 > nImageGened )
            {
                BOOL bNeedCrop = FALSE;
                RECT rectCrop;
                rectCrop.left = rectCrop.top = 0;
                rectCrop.bottom = image.GetWidth();
                rectCrop.right  = image.GetHeight();

                if( (int)image.GetWidth() > nOutWidth )
                {
                    rectCrop.left = ( image.GetWidth() - nOutWidth ) / 2;
                    rectCrop.right = image.GetWidth() / 2 + nOutWidth / 2;
                    bNeedCrop = TRUE;
                }

                if( (int)image.GetHeight() > nOutHeight )
                {
                    rectCrop.top = ( image.GetHeight() - nOutHeight ) / 2;
                    rectCrop.bottom = image.GetHeight() / 2 + nOutHeight / 2;
                    bNeedCrop = TRUE;
                }

                TCHAR szNumberOutputFile[MAX_PATH];
                if( nImageGened > 0 )
                {
                    TCHAR *lpszDot = _tcschr( lpszOutFile, _T('.') );
                    if( lpszDot != NULL )
                    {
                        lpszDot[0] = '\0';
                        _stprintf( szNumberOutputFile, 
                                    _T("%s%d.%s"),
                                    lpszOutFile,
                                    nImageGened,
                                    lpszDot+1 );
                        lpszDot[0] = '.';

                    }
                }
                else
                {
                    _tcscpy( szNumberOutputFile, lpszOutFile );
                }
                //_stprintf( szNumberOutputFile, _T("e:\\jpg\\%f_%d.jpg"), fCurDifSquare, nImageGened );
                    

                if( bNeedCrop )
                {
                    CxImage image2;
                    image.Crop( rectCrop, &image2 );
                    image2.SetJpegQuality(99);
                    image2.Save( szNumberOutputFile ,CXIMAGE_FORMAT_JPG);
                }
                else
                {
                    image.SetJpegQuality(99);
                    image.Save( szNumberOutputFile ,CXIMAGE_FORMAT_JPG );
                }

                if( nImageGened + 1 <  nImageCounter )
                {
                    nImageGened++;
                }
                fPrevlDiffSquare = fCurDifSquare;
            }
            

        }


        delete [] pbBMP;
        
        return 0;
    }

    float GetGrayPixelSize( CxImage &image )
    {
        int nWidth;
        int nHeight;
        nWidth  = (int)image.GetWidth( );
        nHeight = (int)image.GetHeight( );


        int distribution[256];
        ZeroMemory( distribution, sizeof(distribution) );

        float fDiffSquare = 0.0;

        for( int nI = 0; nI < nWidth; nI++ )
        {
            for( int nJ = 0; nJ < nHeight; nJ++ )
            {
                BYTE bVal = image.GetPixelGray( nI, nJ );
                distribution[bVal]++;
            }
        }

        int nVal = nWidth * nHeight / 256;
        float fBase = (float)( nVal * nVal );
        for( int nI = 0; nI < 255; nI++ )
        {
            fDiffSquare += (float)(distribution[nI] - nVal) * (float)(distribution[nI] - nVal)
                            / fBase;
        }

        return fDiffSquare;
    }

    float fPrevlDiffSquare;
};


//
// Main program code
//
int _tmain(int argc, TCHAR* argv[] )
{
    // args: inputfile width height outputfile
    TCHAR szInputFile[MAX_PATH], szOutputFile[MAX_PATH];
    int nWidth;
    int nHeight;
    int nImageCounter = 1;

    if( argc < 5 || !argv )
    {
        _tprintf( TEXT("GrabBitmaps: Param ERROR!\r\n") );
        _tprintf( TEXT("Usage: GrabBitmaps InputFile Width Height OutputFile\r\n"));
        return 0;
    }

    if( argc == 6 )
    {
        nImageCounter = _ttoi( argv[5] );
    }

    
    nWidth    = _ttoi( argv[2] );
    nHeight   = _ttoi( argv[3] );


    if( nWidth == 0 )
    {
        nWidth = 320;
    }
    if( nHeight == 0 )
    {
        nHeight = 240;
    }

    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);

    // Read the filename from the command line      
    _tcsncpy(szInputFile, argv[1], MAX_PATH-1);
    szInputFile[MAX_PATH-1] = 0;      // Null-terminate
    
    // Read the filename from the command line      
    _tcsncpy(szOutputFile, argv[4], MAX_PATH-1);
    szOutputFile[MAX_PATH-1] = 0;      // Null-terminate

    int nSuccess = GrabBitmaps( szInputFile,szOutputFile, nWidth, nHeight, nImageCounter );

    CoUninitialize();

    return nSuccess;
}


int GrabBitmaps(TCHAR * lpszInputFile, TCHAR *lpszOutputFile, int nOutWidth, int nOutHeight, int nImageCounter )
{
    USES_CONVERSION;
    CComPtr< ISampleGrabber > pGrabber;
    CComPtr< IBaseFilter >    pSource;
    CComPtr< IGraphBuilder >  pGraph;
    CComPtr< IVideoWindow >   pVideoWindow;
    HRESULT hr;

    if (!lpszInputFile)
        return -1;

    _tprintf(TEXT("Grabbing bitmaps from %s.\r\n"), lpszInputFile);

    // Create the sample grabber
    //
    pGrabber.CoCreateInstance( CLSID_SampleGrabber );
    if( !pGrabber )
    {
        _tprintf( TEXT("Could not create CLSID_SampleGrabber\r\n") );
        return -1;
    }
    CComQIPtr< IBaseFilter, &IID_IBaseFilter > pGrabberBase( pGrabber );

    // Create the file reader
    //
    pSource.CoCreateInstance( CLSID_AsyncReader );
    if( !pSource )
    {
        _tprintf( TEXT("Could not create source filter\r\n") );
        return -1;
    }

    // Create the graph
    //
    pGraph.CoCreateInstance( CLSID_FilterGraph );
    if( !pGraph )
    {
        _tprintf( TEXT("Could not not create the graph\r\n") );
        return -1;
    }

    // Put them in the graph
    //
    hr = pGraph->AddFilter( pSource, L"Source" );
    hr = pGraph->AddFilter( pGrabberBase, L"Grabber" );

    // Load the source
    //
    CComQIPtr< IFileSourceFilter, &IID_IFileSourceFilter > pLoad( pSource );
    hr = pLoad->Load( T2W( lpszInputFile ), NULL );
    if( FAILED( hr ) )
    {
        _tprintf( TEXT("Could not load the media file\r\n") );
        return -1;
    }

    // Tell the grabber to grab 24-bit video. Must do this
    // before connecting it
    //
    CMediaType GrabType;
    GrabType.SetType( &MEDIATYPE_Video );
    GrabType.SetSubtype( &MEDIASUBTYPE_RGB24 );
    hr = pGrabber->SetMediaType( &GrabType );

    // Get the output pin and the input pin
    //
    CComPtr< IPin > pSourcePin;
    CComPtr< IPin > pGrabPin;

    pSourcePin = GetOutPin( pSource, 0 );
    pGrabPin   = GetInPin( pGrabberBase, 0 );


    // ... and connect them
    //
    hr = pGraph->Connect( pSourcePin, pGrabPin );
    if( FAILED( hr ) )
    {
        _tprintf( TEXT("Could not connect source filter to grabber\r\n") );
        return -1;
    }

    // This semi-COM object will receive sample callbacks for us
    //
    CSampleGrabberCB CB( lpszOutputFile, nOutWidth, nOutHeight, nImageCounter );

    // Ask for the connection media type so we know its size
    //
    AM_MEDIA_TYPE mt;
    hr = pGrabber->GetConnectedMediaType( &mt );

    VIDEOINFOHEADER * vih = (VIDEOINFOHEADER*) mt.pbFormat;
    CB.Width  = vih->bmiHeader.biWidth;
    CB.Height = vih->bmiHeader.biHeight;
    FreeMediaType( mt );

    // Render the grabber output pin (to a video renderer)
    //
    CComPtr <IPin> pGrabOutPin = GetOutPin( pGrabberBase, 0 );
    hr = pGraph->Render( pGrabOutPin );
    if( FAILED( hr ) )
    {
        _tprintf( TEXT("Could not render grabber output pin\r\n") );
        return -1;
    }

    // Don't buffer the samples as they pass through
    //
    hr = pGrabber->SetBufferSamples( FALSE );

    // Only grab one at a time, stop stream after
    // grabbing one sample
    //
    hr = pGrabber->SetOneShot( TRUE );

    // Set the callback, so we can grab the one sample
    //
    hr = pGrabber->SetCallback( &CB, 1 );

    // Get the seeking interface, so we can seek to a location
    //
    CComQIPtr< IMediaSeeking, &IID_IMediaSeeking > pSeeking( pGraph );

    // Query the graph for the IVideoWindow interface and use it to
    // disable AutoShow.  This will prevent the ActiveMovie window from
    // being displayed while we grab bitmaps from the running movie.
    CComQIPtr< IVideoWindow, &IID_IVideoWindow > pWindow = pGraph;
    if (pWindow)
    {
        hr = pWindow->put_AutoShow(OAFALSE);
    }

    srand( (int)time(NULL) );

    LONGLONG llDuration = 0;
    pSeeking->GetDuration( &llDuration );
    REFERENCE_TIME Start = 0;

    // Find a limited number of frames
    //
    int nImages = nImageCounter < NUM_FRAMES_TO_GRAB ? NUM_FRAMES_TO_GRAB: nImageCounter * 5;
    Start = (llDuration / nImages) / 20 * ( rand() % 20 );
    for( int i = 0; i < nImages - 1; i++ )
    {
        // set position
        Start += ( llDuration / nImages );
        
        hr = pSeeking->SetPositions( &Start, 
                                     AM_SEEKING_AbsolutePositioning | AM_SEEKING_SeekToKeyFrame, 
                                     NULL, 
                                     AM_SEEKING_NoPositioning );

        // activate the threads
        CComQIPtr< IMediaControl, &IID_IMediaControl > pControl( pGraph );
        hr = pControl->Run( );

        // wait for the graph to settle
        CComQIPtr< IMediaEvent, &IID_IMediaEvent > pEvent( pGraph );
        long EvCode = 0;

        hr = pEvent->WaitForCompletion( 30 * 1000, &EvCode );
        if( FAILED( hr ) )
        {
            exit( 0 );
        }
        // callback wrote the sample
    }

    _tprintf(TEXT("Sample grabbing complete.\r\n"));
    return 0;
}


HRESULT GetPin( IBaseFilter * pFilter, PIN_DIRECTION dirrequired, int iNum, IPin **ppPin)
{
    CComPtr< IEnumPins > pEnum;
    *ppPin = NULL;

    HRESULT hr = pFilter->EnumPins(&pEnum);
    if(FAILED(hr)) 
        return hr;

    ULONG ulFound;
    IPin *pPin;
    hr = E_FAIL;

    while(S_OK == pEnum->Next(1, &pPin, &ulFound))
    {
        PIN_DIRECTION pindir = (PIN_DIRECTION)3;

        pPin->QueryDirection(&pindir);
        if(pindir == dirrequired)
        {
            if(iNum == 0)
            {
                *ppPin = pPin;  // Return the pin's interface
                hr = S_OK;      // Found requested pin, so clear error
                break;
            }
            iNum--;
        } 

        pPin->Release();
    } 

    return hr;
}


IPin * GetInPin( IBaseFilter * pFilter, int nPin )
{
    CComPtr<IPin> pComPin=0;
    GetPin(pFilter, PINDIR_INPUT, nPin, &pComPin);
    return pComPin;
}


IPin * GetOutPin( IBaseFilter * pFilter, int nPin )
{
    CComPtr<IPin> pComPin=0;
    GetPin(pFilter, PINDIR_OUTPUT, nPin, &pComPin);
    return pComPin;
}




